#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import re
import csv
import time
import signal
import statistics
import subprocess
from datetime import datetime
from influxdb_client import Point

from handle_influxdb import InfluxDBHandler

do_influxdb = InfluxDBHandler()
# ============= 配置参数 =============
# 要监控的进程名(支持模糊匹配)
PROCESS_NAME = "trackdemo"
# 监控间隔(秒)
INTERVAL = 2
# 输出CSV文件基础名称（将自动添加时间戳）
OUTPUT_FILE_BASE = "process_monitor"
# ===================================

# ANSI颜色控制字符的正则表达式
ANSI_ESCAPE = re.compile(r'\x1b\[[\d;]*[a-zA-Z]')

# 全局变量用于存储CPU和内存使用数据
cpu_values = []
mem_values = []


def get_process_info(process_name):
    """
    查找匹配进程名的进程，并返回PID和CPU使用率
    基于Android设备特定的top命令输出格式
    """
    try:
        # 使用正则表达式技巧避免grep自身被匹配到 - 将第一个字母放入[]中
        first_char = process_name[0] if process_name else ""
        pattern = f"{process_name}" if not first_char else f"[{first_char}]{process_name[1:]}"
        cmd = f"adb shell \"top -n 1 | grep {pattern}\""
        result = subprocess.check_output(cmd, shell=True, text=True)

        if not result.strip():
            print(f"未找到匹配 '{process_name}' 的进程")
            return None, None, None

        # 取第一行（如果有多行结果）
        process_line = result.strip().split('\n')[0]
        print(f"找到进程行: [{process_line}]")

        # 移除ANSI颜色控制字符
        process_line = ANSI_ESCAPE.sub('', process_line)
        print(f"移除ANSI控制字符后: [{process_line}]")

        # 处理行首可能存在的多余空格
        process_line = process_line.strip()

        # 按空格分割字段
        fields = re.split(r'\s+', process_line)

        print(f"分割后的字段: {fields}")

        if len(fields) < 10:
            print(f"行格式异常，无法解析: {process_line}")
            return None, None, None

        # 第一个字段就是PID
        pid = fields[0]

        # 尝试查找CPU使用率
        # 在Android设备的top输出中CPU使用率通常在R或S状态标记后面一列
        cpu_usage = None
        for i, field in enumerate(fields):
            if field == "R" or field == "S":
                if i + 1 < len(fields):
                    cpu_usage = fields[i + 1]
                    break

        # 如果上面的方法失败，尝试直接使用[%CPU]列
        if not cpu_usage and len(fields) > 8:
            # 通常在第8或9列
            for i in range(7, min(10, len(fields))):
                if fields[i].replace(".", "").isdigit() or fields[i].endswith("%"):
                    cpu_usage = fields[i].replace("%", "")
                    break

        # 找不到CPU使用率，使用备用位置
        if not cpu_usage:
            cpu_usage = "未知"

        # 查找进程名的位置
        process_idx = -1
        for i, field in enumerate(fields):
            if process_name in field:
                process_idx = i
                break

        # 提取进程名和参数
        if process_idx >= 0:
            process_full_name = " ".join(fields[process_idx:])
        else:
            process_full_name = process_name

        print(f"解析结果 - PID: {pid}, CPU: {cpu_usage}, 进程: {process_full_name}")
        return pid, cpu_usage, process_full_name

    except subprocess.CalledProcessError as e:
        print(f"执行命令失败: {e}")
        return None, None, None
    except Exception as e:
        print(f"获取进程信息时出错: {e}")
        return None, None, None


def get_memory_info(pid):
    """
    获取指定PID的内存使用情况
    """
    try:
        if not pid:
            print("无法获取内存信息：PID为空")
            return None

        result = subprocess.check_output(f"adb shell dumpsys meminfo {pid}", shell=True, text=True)

        # 查找TOTAL PSS行
        match = re.search(r'TOTAL PSS:\s+(\d+)', result)
        if match:
            mem_usage = match.group(1)
            print(f"内存使用: {mem_usage}KB")
            return mem_usage

        print("未找到内存信息")
        return None
    except Exception as e:
        print(f"获取内存信息时出错: {e}")
        return None


def print_statistics():
    """
    打印CPU和内存使用的统计信息
    """
    print("\n===== 监控统计信息 =====")

    print("\nCPU使用率(%)统计:")
    if cpu_values:
        # 将所有字符串转换为浮点数（过滤掉非数字值）
        numeric_cpu = [float(v) for v in cpu_values if v.replace('.', '', 1).isdigit()]
        if numeric_cpu:
            print(f"  - 最大值: {max(numeric_cpu):.1f}%")
            print(f"  - 最小值: {min(numeric_cpu):.1f}%")
            print(f"  - 平均值: {statistics.mean(numeric_cpu):.1f}%")
            print(f"  - 样本数: {len(numeric_cpu)}")
        else:
            print("  没有有效的CPU数据")
    else:
        print("  没有收集到CPU数据")

    print("\n内存使用(KB)统计:")
    if mem_values:
        # 将所有字符串转换为整数（过滤掉非数字值）
        numeric_mem = [int(v) for v in mem_values if v.isdigit()]
        if numeric_mem:
            print(f"  - 最大值: {max(numeric_mem):,} KB")
            print(f"  - 最小值: {min(numeric_mem):,} KB")
            print(f"  - 平均值: {statistics.mean(numeric_mem):,.1f} KB")
            print(f"  - 样本数: {len(numeric_mem)}")
        else:
            print("  没有有效的内存数据")
    else:
        print("  没有收集到内存数据")


def signal_handler(sig, frame):
    """
    处理Ctrl+C信号，在退出前打印统计数据
    """
    print("\n\n正在退出监控...")
    print_statistics()
    print("\n监控已停止")
    exit(0)


def main():
    # 注册信号处理器以捕获Ctrl+C
    signal.signal(signal.SIGINT, signal_handler)

    # 检查adb是否可用
    try:
        devices = subprocess.check_output("adb devices", shell=True, text=True)
        if "device" not in devices:
            print("错误: 未检测到已连接的设备")
            return
        print(f"已连接设备: {devices}")
    except Exception as e:
        print(f"错误: 无法连接到adb: {e}")
        return

    # 生成带时间戳的输出文件名
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    output_file = f"{OUTPUT_FILE_BASE}_{timestamp}.csv"

    # 创建CSV文件
    with open(output_file, 'w', newline='') as csvfile:
        fieldnames = ['时间戳', 'PID', '进程名', 'CPU使用率(%)', '内存使用(KB)']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()

    print(f"开始监控进程 '{PROCESS_NAME}'，数据将写入 {output_file}")
    print(f"监控间隔: {INTERVAL}秒, 按Ctrl+C停止")

    try:
        while True:
            timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            print(f"\n===== {timestamp} 开始新的监控周期 =====")

            # 获取进程信息
            pid, cpu_usage, process_full_name = get_process_info(PROCESS_NAME)

            if pid:
                # 获取内存信息
                mem_usage = get_memory_info(pid)

                # 记录数据用于统计
                if cpu_usage and cpu_usage != "未知":
                    cpu_values.append(cpu_usage)

                if mem_usage:
                    mem_values.append(mem_usage)
                if cpu_usage and cpu_usage != "未知" and mem_usage:
                    measurements = [
                        Point("device_001").tag("pid", "main_pid").field("cpu", float(cpu_usage)),
                        Point("device_001").tag("pid", "main_pid").field("mem", float(mem_usage)),
                    ]
                    do_influxdb.write_mul_data(measurements=measurements)

                # 写入CSV
                with open(output_file, 'a', newline='') as csvfile:
                    writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                    row_data = {
                        '时间戳': timestamp,
                        'PID': pid,
                        '进程名': process_full_name,
                        'CPU使用率(%)': cpu_usage,
                        '内存使用(KB)': mem_usage if mem_usage else "N/A"
                    }
                    writer.writerow(row_data)

                print(
                    f"监控结果: PID: {pid}, 进程: {process_full_name}, CPU: {cpu_usage}%, 内存: {mem_usage if mem_usage else 'N/A'}KB")
            else:
                print(f"找不到匹配 '{PROCESS_NAME}' 的进程")

            # 等待下一次监控
            print(f"等待 {INTERVAL} 秒...")
            time.sleep(INTERVAL)

    except KeyboardInterrupt:
        # 捕获Ctrl+C，优雅退出
        print("\n\n正在退出监控...")
        print_statistics()
        print("\n监控已停止")


if __name__ == "__main__":
    main()